import json
import csv
import os


def convert_jsonl_to_csv(input_file, output_file):
    all_data = []
    with open(input_file, 'r', encoding='utf-8') as f:
        for line in f:
            line = line.strip()
            data = json.loads(line)
            all_data.append(data)

    fieldnames = set()
    for item in all_data:
        fieldnames.update(item.keys())
    fieldnames = sorted(list(fieldnames))

    # 先按generated降序排列，再按prompt_id升序排列
    all_data.sort(key=lambda x: (-float(x.get('generated', 0) or 0), str(x.get('prompt_id', ''))))

    with open(output_file, 'w', newline='', encoding='utf-8') as f:
        writer = csv.DictWriter(f, fieldnames=fieldnames)
        writer.writeheader()
        writer.writerows(all_data)

    print(f"已将 {len(all_data)} 条记录从 {os.path.basename(input_file)} 转换到 {os.path.basename(output_file)}")


def main():
    input_dir = './data'
    input_files = [
        os.path.join(input_dir, 'train.jsonl'),
        os.path.join(input_dir, 'train_prompts.jsonl'),
        os.path.join(input_dir, 'test.jsonl')
    ]

    for input_file in input_files:
        if not os.path.exists(input_file):
            print(f"文件不存在: {input_file}")
            continue

        base_name = os.path.splitext(os.path.basename(input_file))[0]
        output_file = os.path.join(input_dir, f"{base_name}_sorted.csv")

        print(f"正在处理: {os.path.basename(input_file)} -> {os.path.basename(output_file)}")
        convert_jsonl_to_csv(input_file, output_file)


if __name__ == "__main__":
    main()
